#include "kernel_operator.h"
using namespace AscendC;

template<typename T>
class KernelScatterSub {
 public:
  __aicore__ inline KernelScatterSub() {}
  __aicore__ inline void Init(GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR var_ref,
                              bool use_locking, uint32_t size, uint32_t indice_length, uint32_t datatype) {

    this->datatype = datatype;
    this->use_locking = use_locking;
    this->origin_total_size = size;
    this->origin_indice_length = indice_length;
    this->total_size = ((size+this->block_size - 1)/this->block_size) * this->block_size;
    this->indice_length = ((indice_length+this->block_size - 1)/this->block_size) * this->block_size;

    varGm.SetGlobalBuffer((__gm__ T*)var, this->total_size);
    indicesGm.SetGlobalBuffer((__gm__ int32_t *)indices, this->indice_length);
    updatesGm.SetGlobalBuffer((__gm__ T*)updates, this->indice_length);
    yGm.SetGlobalBuffer((__gm__ T*)var_ref, this->total_size);
    
    pipe.InitBuffer(inQueueIndices, 1, this->indice_length * sizeof(int32_t));
    pipe.InitBuffer(inQueueUpdates, 1, this->indice_length * sizeof(T));
    pipe.InitBuffer(inQueueVar, 1, this->total_size * sizeof(T));
    pipe.InitBuffer(outQueueOUT, 1, this->total_size * sizeof(T));
  }
  __aicore__ inline void Process() {
    CopyIn();
    Compute();
    CopyOut();
  }

 private:
  __aicore__ inline void CopyIn() {
    LocalTensor<T> varLocal = inQueueVar.AllocTensor<T>();
    LocalTensor<T> updatesLocal = inQueueUpdates.AllocTensor<T>();
    LocalTensor<int32_t> indicesLocal = inQueueIndices.AllocTensor<int32_t>();
    //搬运varLocal
    if(this->origin_total_size >=this->block_size ){
        if(this->origin_total_size == this->total_size){
            DataCopy(varLocal, varGm,this->total_size);
        }else{
            int loop = this->origin_total_size / this->block_size;
            for(int i=0; i<loop; i++){
                DataCopy(varLocal[i * this->block_size], varGm,this->block_size);
            }
            int left = this->origin_total_size % this->block_size;
            for(int i=0; i<left; i++){
                varLocal.SetValue(loop * this->block_size + i, varGm.GetValue(loop*this->block_size + i));
            }
            for(int i=this->origin_total_size; i<this->block_size; i++){
                varLocal.SetValue(loop * this->block_size + i, (T)0);
            }
        }   
    }else{
        for(int i=0; i<this->origin_total_size; i++){
            varLocal.SetValue(i, varGm.GetValue(i));
        }
        for(int i=origin_total_size; i<this->total_size; i++){
            varLocal.SetValue(i, (T)0);
        }
    }
    //搬运updatesLocal,indicesLocal
    if(this->origin_indice_length >=this->block_size ){
        if(this->origin_indice_length == this->indice_length){
            DataCopy(updatesLocal, updatesGm,this->indice_length);
            DataCopy(indicesLocal, indicesGm,this->indice_length);
        }else{
            int loop = this->origin_indice_length / this->block_size;
            for(int i=0; i<loop; i++){
                DataCopy(updatesLocal, updatesGm,this->block_size);
                DataCopy(indicesLocal, indicesGm,this->block_size);
            }
            int left = this->origin_indice_length % this->block_size;
            for(int i=0; i<left; i++){
                updatesLocal.SetValue(loop * this->block_size + i, updatesGm.GetValue(loop*this->block_size + i));
                indicesLocal.SetValue(loop * this->block_size + i, indicesGm.GetValue(loop*this->block_size + i));
            }
            for(int i=this->origin_indice_length; i<this->block_size; i++){
                updatesLocal.SetValue(loop * this->block_size + i, (T)0);
                indicesLocal.SetValue(loop * this->block_size + i, -1);
            }
        }
    }else{
        for(int i=0; i<this->origin_indice_length; i++){
            updatesLocal.SetValue(i, updatesGm.GetValue(i));
            indicesLocal.SetValue(i, indicesGm.GetValue(i));
        }
        for(int i=origin_indice_length; i<this->indice_length; i++){
            updatesLocal.SetValue(i, (T)0);
            indicesLocal.SetValue(i, -1);
        }
    }

    // DataCopy(updatesLocal, updatesGm,this->indice_length);
    // DataCopy(indicesLocal, indicesGm,this->indice_length);

    inQueueVar.EnQue(varLocal);
    inQueueUpdates.EnQue(updatesLocal);
    inQueueIndices.EnQue(indicesLocal);
  }
    
  __aicore__ inline void Compute() {
    LocalTensor<T> varLocal = inQueueVar.DeQue<T>();
    LocalTensor<T> updatesLocal = inQueueUpdates.DeQue<T>();
    LocalTensor<int32_t> indicesLocal = inQueueIndices.DeQue<int32_t>();

    LocalTensor<T> yLocal = outQueueOUT.AllocTensor<T>();
    pipe_barrier(PIPE_ALL);
    int location;
    if(this->datatype == 0 || this->datatype == 1){
        for(int i = 0; i<this->origin_indice_length; i++){
            location = indicesLocal.GetValue(i);
            if(location != -1)
                varLocal.SetValue(location,(float)varLocal.GetValue(location) - (float)updatesLocal.GetValue(i));
        }   
    }else{
        for(int i = 0; i<this->origin_indice_length; i++){
            location = indicesLocal.GetValue(i);
            if(location != -1)
                varLocal.SetValue(location,varLocal.GetValue(location) - updatesLocal.GetValue(i));
        } 
    }
    pipe_barrier(PIPE_ALL);


    DataCopy(yLocal, varLocal, this->total_size);
    
    outQueueOUT.EnQue<T>(yLocal);

    inQueueVar.FreeTensor(varLocal);
    inQueueUpdates.FreeTensor(updatesLocal);
    inQueueIndices.FreeTensor(indicesLocal);
  }
  __aicore__ inline void CopyOut() {
    LocalTensor<T> yLocal = outQueueOUT.DeQue<T>();
    
    // if(this->origin_total_size >=this->block_size ){
    //     if(this->origin_total_size == this->total_size){
    DataCopy(yGm, yLocal, this->total_size);
    DataCopy(varGm, yLocal, this->total_size);

    //     }else{
    //         int loop = this->origin_total_size / this->block_size;
    //         for(int i=0; i<loop; i++){
    //             DataCopy(yGm[i*this->block_size], yLocal[i * this->block_size],this->block_size);
    //         }
    //         int left = this->origin_total_size % this->block_size;
    //         for(int i=0; i<left; i++){
    //             yGm.SetValue(loop * this->block_size + i, yLocal.GetValue(loop*this->block_size + i));
    //         }
    //     }   
    // }else{
    //     for(int i=0; i<this->origin_total_size; i++){
    //         yGm.SetValue(i, yLocal.GetValue(i));
    //     }
    // }
    

    outQueueOUT.FreeTensor(yLocal);
  }

 private:
  TPipe pipe;
  // TQue<QuePosition::VECIN, 1> inQueueX, inQueueY, inQueueZ;
  TQue<QuePosition::VECIN, 1> inQueueVar;
  TQue<QuePosition::VECIN, 1> inQueueIndices, inQueueUpdates;
  TQue<QuePosition::VECOUT, 1> outQueueOUT;
  GlobalTensor<T> varGm;
  GlobalTensor<int32_t> indicesGm;
  GlobalTensor<T> updatesGm;
  GlobalTensor<T> yGm;
  bool use_locking;
  uint32_t total_size;
  uint32_t origin_total_size;
  uint32_t indice_length;
  uint32_t origin_indice_length;
  uint32_t block_size = 32;
  uint32_t datatype;
};
extern "C" __global__ __aicore__ void scatter_sub(GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR var_ref, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    if(tiling_data.datatype == 1){
        //half
        KernelScatterSub<half> op;
        op.Init(var, indices, updates, var_ref, tiling_data.use_locking, tiling_data.size, tiling_data.indice_length, tiling_data.datatype);
    }else if(tiling_data.datatype == 0){
        //float32
        KernelScatterSub<float> op;
        op.Init(var, indices, updates, var_ref, tiling_data.use_locking, tiling_data.size, tiling_data.indice_length, tiling_data.datatype);
    }else if(tiling_data.datatype == 2){
        //int8
        KernelScatterSub<int8_t> op;
        op.Init(var, indices, updates, var_ref, tiling_data.use_locking, tiling_data.size, tiling_data.indice_length, tiling_data.datatype);
    }else if(tiling_data.datatype == 3){
        //int32
        KernelScatterSub<int32_t> op;
        op.Init(var, indices, updates, var_ref, tiling_data.use_locking, tiling_data.size, tiling_data.indice_length, tiling_data.datatype);
    }
}